package dfxnftables

import (
	"context"
	"fmt"
	"net"

	"github.com/google/nftables"
)
// NfSet 表示一个 nftables 集合的结构体，用于管理 IP 地址集合。
type NfSet struct {
	NfMeta
	SetIPList []string // 要添加到集合中的 IP 列表
	SetName   string   // 集合名称
}

// CreateSet 在指定的 nftables 表中创建一个新的集合。
// 参数:
//   - ctx: 上下文，用于控制请求的生命周期。
//   - c: 与 nftables 通信的连接对象。
// 返回值:
//   - error: 如果创建失败则返回错误信息。
func (n *NfSet) CreateSet(ctx context.Context, c *nftables.Conn) error {
	if err := n.NfMeta.Validate(); err != nil {
		return err
	}
	tbl, err := n.ResolveTable(c)
	if err != nil {
		return err
	}
	fam := n.Family()
	// 定义集合 key 类型
	keyType := nftables.TypeIPAddr
	if fam == nftables.TableFamilyIPv6 {
		keyType = nftables.TypeIP6Addr
	}
	s := &nftables.Set{
		Table:   tbl,
		Name:    n.SetName,
		KeyType: keyType,
	}
	if err := c.AddSet(s, nil); err != nil {
		return fmt.Errorf("add set failed: %w", err)
	}
	return c.Flush()
}

// CheckSet 检查指定名称的集合是否存在于对应的 nftables 表中。
// 参数:
//   - ctx: 上下文，用于控制请求的生命周期。
//   - c: 与 nftables 通信的连接对象。
// 返回值:
//   - bool: 如果存在返回 true，否则返回 false。
//   - error: 如果查询过程中发生错误，则返回错误信息。
func (n *NfSet) CheckSet(ctx context.Context, c *nftables.Conn) (bool, error) {
	if err := n.NfMeta.Validate(); err != nil {
		return false, err
	}
	tbl, err := n.ResolveTable(c)
	if err != nil {
		return false, err
	}
	s, err := c.GetSetByName(tbl, n.SetName)
	if err != nil {
		return false, err
	}
	return s != nil, nil
}

// QuerySet 查询指定集合中的所有元素，并将其转换为字符串格式的 IP 地址列表。
// 参数:
//   - ctx: 上下文，用于控制请求的生命周期。
//   - c: 与 nftables 通信的连接对象。
// 返回值:
//   - []string: 包含集合中所有 IP 地址的字符串切片。
//   - error: 如果查询失败则返回错误信息。
func (n *NfSet) QuerySet(ctx context.Context, c *nftables.Conn) ([]string, error) {
	if err := n.NfMeta.Validate(); err != nil {
		return nil, err
	}
	fam := n.Family()
	tbl, err := n.ResolveTable(c)
	if err != nil {
		return nil, err
	}
	s, err := c.GetSetByName(tbl, n.SetName)
	if err != nil {
		return nil, err
	}
	if s == nil {
		return nil, fmt.Errorf("set %s not found", n.SetName)
	}
	els, err := c.GetSetElements(s)
	if err != nil {
		return nil, err
	}
	var out []string
	for _, e := range els {
		if fam == nftables.TableFamilyIPv6 {
			if len(e.Key) >= 16 {
				ip := net.IP(e.Key[:16])
				out = append(out, ip.String())
			}
		} else {
			if len(e.Key) >= 4 {
				ip4 := net.IPv4(e.Key[0], e.Key[1], e.Key[2], e.Key[3])
				out = append(out, ip4.String())
			}
		}
	}
	return out, nil
}

// DeleteSet 删除指定名称的集合。
// 参数:
//   - ctx: 上下文，用于控制请求的生命周期。
//   - c: 与 nftables 通信的连接对象。
// 返回值:
//   - error: 如果删除失败则返回错误信息。
func (n *NfSet) DeleteSet(ctx context.Context, c *nftables.Conn) error {
	if err := n.NfMeta.Validate(); err != nil {
		return err
	}
	tbl, err := n.ResolveTable(c)
	if err != nil {
		return err
	}
	s, err := c.GetSetByName(tbl, n.SetName)
	if err != nil {
		return err
	}
	if s == nil {
		return fmt.Errorf("set %s not found", n.SetName)
	}
	c.DelSet(s)
	return c.Flush()
}

// FlushSet 清空指定集合中的所有元素。
// 参数:
//   - ctx: 上下文，用于控制请求的生命周期。
//   - c: 与 nftables 通信的连接对象。
// 返回值:
//   - error: 如果清空失败则返回错误信息。
func (n *NfSet) FlushSet(ctx context.Context, c *nftables.Conn) error {
	if err := n.NfMeta.Validate(); err != nil {
		return err
	}
	tbl, err := n.ResolveTable(c)
	if err != nil {
		return err
	}
	s, err := c.GetSetByName(tbl, n.SetName)
	if err != nil {
		return err
	}
	if s == nil {
		return fmt.Errorf("set %s not found", n.SetName)
	}
	els, err := c.GetSetElements(s)
	if err != nil {
		return err
	}
	if len(els) == 0 {
		return nil
	}
	if err := c.SetDeleteElements(s, els); err != nil {
		return err
	}
	return c.Flush()
}